#ifndef CLASS_DISK
#define CLASS_DISK

#include <stdio.h>

#include "P.h"
#include "AG.h"
#include "TreeForceFinder.h"
#include "Vector.h"


#define SQR(x) ((x)*(x))

typedef unsigned int uint;
typedef long unsigned int luint;

using namespace std;


class Disk {
	public:
		
		int n;
		double m;
		double r0;
		double z0;
		double rmax;
		double zmax;

		double G;
	
		//derivatives
		P *p;
		
		Disk () :
		n(0),
		m(0),
		r0(2.5),
		z0(0.5),
		rmax(100),
		zmax(100),
		G(1),
		p(0)
		{}
		
		~Disk () {
			free();
		}
		
		void set() {
			
			free();
			p = new P[n];
			
			set_pos(p, n);
			set_vel_uniform(p, n);
			
		}
		
		
		double rho(double r, double z) {
			return m/(4*M_PI*z0*r0*r0)*(1-pow(tanh(z/z0),2))*exp(-r/r0);
		};

		double M(double r1, double z1, double r2, double z2, int nr=1, int nz=1) {
			
			double m(0);
			
			double dr = (r2-r1)/nr;
			double dz = (z2-z1)/nz;
			
			for (int iz=0; iz< nz; iz++)
			for (int ir=0; ir< nr; ir++) {
				
				double rci = r1 + dr*(ir+0.5);
				double zci = z1 + dz*(iz+0.5);
				
				//double rzi = sqrt(SQR(rci) + SQR(zci));
				
				m+=rho(rci, zci)*rci;
			}
			
			return 4*M_PI*m*dr*dz;

		}

		
		void set_pos(P *p, int n) {
			
			printf("Disk::set_pos()... "); fflush(stdout);
		
			srand48(222);

			double r, z, r_old;
			
			for(int i = 0; i<n; i++)  {
			
				do {
				
					double q = drand48();
					z = z0 / 2 * log(q / (1 - q));
					q = drand48();
					r = 1.0;
					
					do	{
					
						double f = (1 + r) * exp(-r) + q - 1;
						double f_ = -r * exp(-r);

						r_old = r;
						r = r - f / f_;
					
					} while(fabs(r - r_old) / r > 1e-7);

					r *= r0;

				} while (r > rmax || fabs(z) > zmax);	//countr++; else i++;
				
				double phi = drand48() * M_PI* 2;
				
				p[i].Pos[0] = r * cos(phi);
				p[i].Pos[1] = r * sin(phi);
				p[i].Pos[2] = z;
				
			}
		
			for(int i=0; i<n; i++) p[i].Mass = m/n;
			
			printf("done\n");
			
		}

		void set_pot(P *p, int n) {

			TreeForceFinder tff;
			tff.G = G;
			tff.NumPart = n;
			tff.force_treeallocate(1.5*n, n);
			tff.force_treebuild(p);
			
			REAL acc[3];
			for(int i=0; i<n; i++)
				tff.force_treeevaluate(p, p[i].Pos, acc, p[i].pot);
			
			tff.force_treefree();
				
		}

		void set_pot(AG &ag) {
			for(int i=0; i<n; i++) { 
				AG::A a = ag.get_acc_pot(p[i].Pos);
				p[i].pot = a.pot;
			}
		}		
		
		void set_vel_uniform(P *p, int n) {
			
			printf("Disk::set_vel_uniform()...\n"); fflush(stdout);
			
			TreeForceFinder tff;
			tff.G = G;
			tff.NumPart = n; 
			tff.force_treeallocate(1.5*n, n);
			tff.force_treebuild(p);
			
			int i(0);
			while ( i<n) {
				
				REAL posi[] = {12,0,2.5}; 
				REAL acci[] = {0,0,0};
				REAL pot(0);
				
				tff.force_treeevaluate(p, p[i].Pos, acci, pot);
				
				double v_esc = sqrt(-pot); //must be sqrt(-2*pot);
				
				p[i].Vel[0] = v_esc*(2*drand48() - 1);
				p[i].Vel[1] = v_esc*(2*drand48() - 1);
				p[i].Vel[2] = v_esc*(2*drand48() - 1);
				p[i].v_esc = v_esc;
				
				
				double v = sqrt(SQR(p[i].Vel[0]) + SQR(p[i].Vel[1]) + SQR(p[i].Vel[2]));
				if ( v_esc < v ) continue; 
				
				i++;
			}
			
			
			tff.force_treefree();
			printf("done\n");
		}
		
		void flip_vels(){
			
			for (int i=0; i<n; i++) {
				
				Vector<REAL> &R_ = *(Vector<REAL>*)p[i].Pos;
				Vector<REAL> &V = *(Vector<REAL>*)p[i].Vel;
				Vector<REAL> R(p[i].Pos[0], p[i].Pos[1], 0);
				Vector<REAL> Z(0,0,1);
				Vector<REAL> T(Z%R);
				
				
				
				Vector<REAL> I(R_%V);
				//Vector<REAL> T(R%I);
				Vector<REAL> VR = (V*R.ort())*R.ort();
				Vector<REAL> VT = (V*T.ort())*T.ort();
				Vector<REAL> VZ = (V*Z.ort())*Z.ort();
				
				
				
				if (V*T.ort()<0) {
					VT*=-1;
					VZ*=-1;
					V = VR+VT+VZ;
					//V*=-1; 
				}
				
				/*
				if (I.z<0) {
					//VT*=-1;
					//VZ*=-1;
					//V = VR+VT+VZ;
					V*=-1; 
				}
				*/
				
				
				/*
				if ((R%V).z<0) {
					printf("warning!!!\n");
					(R%V).print();
					//exit(0);
				}
				
				*/
				
			}
			
			
			
			
		}

		void remove_drift(){
			
			Vector<REAL> V;
			for (int i=0; i<n; i++)
				V+= *(Vector<REAL>*)p[i].Vel;
			
			V/=n;
			
			for (int i=0; i<n; i++)
				*(Vector<REAL>*)p[i].Vel -=V;
		
			//V.print();
		}
		
		
		void set_vel_uniform(P *p, int n, AG &ag) {

			printf("Disk::set_vel_uniform(p, %i, ag)...\n", n); fflush(stdout);
			
			int i(0);
			while ( i<n ) {

				double r = sqrt( SQR(p[i].Pos[0]) + SQR(p[i].Pos[1]) );
				double z = fabs(p[i].Pos[2]);
				
				AG::A a = ag.get_acc_pot(r,z);
				double v_esc = sqrt(-2*a.pot);
				
				p[i].Vel[0] = v_esc*(2*drand48() - 1);
				p[i].Vel[1] = v_esc*(2*drand48() - 1);
				p[i].Vel[2] = v_esc*(2*drand48() - 1);
				p[i].v_esc = v_esc;
				//printf("%g\n", v_esc);
				
				double v = sqrt(SQR(p[i].Vel[0]) + SQR(p[i].Vel[1]) + SQR(p[i].Vel[2]));
				if ( v_esc < v ) continue; 
				
				i++;
			}
			
		
			
			printf("done\n");
			
		}
		
		
		void calc_and_save_w(	const char *fnr,
										const char *fnt,
										const char *fnz ) {
			
			printf("Disk::calc_and_save_w(%s, %s, %s)...", fnr, fnt, fnz); fflush(stdout);

			
			class SG { public:
				uint nr;
				double r1, r2;
				SG(): nr(4),r1(1),r2(15) {};
			}sg;
			
			class VG { public:
				uint nr, nt, nz;
				VG(): nr(30),nt(30),nz(30) {};
			}vg;
			
			uint wnr = sg.nr*vg.nr;
			uint wnt = sg.nr*vg.nt;
			uint wnz = sg.nr*vg.nz;
			
			double *wr = new double[wnr];
			double *wt = new double[wnt];
			double *wz = new double[wnz];
			
			memset(wr, 0, wnr*sizeof(*wr) );
			memset(wt, 0, wnt*sizeof(*wt) );
			memset(wz, 0, wnz*sizeof(*wz) );
			
			
			for (int i=0; i<n; i++) {
				
				double  r = sqrt(SQR(p[i].Pos[0])+SQR(p[i].Pos[1])); 
				double _r = r*(sg.r2-sg.r1)/sg.r2 + sg.r1;
				uint ir = sg.nr*log(_r/sg.r1)/log(sg.r2/sg.r1);
			
				if (sg.nr <= ir) continue;
				
				Vector<REAL> V = *(Vector<REAL>*)p[i].Vel / sqrt(-2*p[i].pot);
				
				Vector<REAL> R(p[i].Pos[0],p[i].Pos[1], 0);
				Vector<REAL> Z(0,0,1);
				Vector<REAL> T(Z%R);
				
				double vr = V*R.ort();
				double vt = V*T.ort();
				double vz = V*Z.ort();
				
				//double rc = (sg.r1*pow(sg.r2/sg.r1,(double)(i+0.5)/sg.nr)-sg.r1)*sg.r2/(sg.r2-sg.r1);
				double v_max  = 1;
				
				double dvr = 2.0/vg.nr;
				double dvt = 2.0/vg.nt;
				double dvz = 2.0/vg.nz;
				
				int vrid = int(floor(vr/dvr) ) + vg.nr/2;
				int vtid = int(floor(vt/dvt) ) + vg.nt/2;
				int vzid = int(floor(vz/dvz) ) + vg.nz/2;
				
				if (0 <= vrid && vrid < vg.nr) wr[ir*vg.nr+vrid]++;
				if (0 <= vtid && vtid < vg.nt) wt[ir*vg.nt+vtid]++;
				if (0 <= vzid && vzid < vg.nz) wz[ir*vg.nz+vzid]++;
				
			}

			for (int i=0; i<sg.nr; i++){
				double wrs(0), wts(0), wzs(0); 
				for (int k=0; k<vg.nr; k++) wrs+= wr[i*vg.nr+k];
				for (int k=0; k<vg.nt; k++) wts+= wt[i*vg.nt+k];
				for (int k=0; k<vg.nz; k++) wzs+= wz[i*vg.nz+k];
				
				for (int k=0; k<vg.nr; k++) wr[i*vg.nr+k]/=wrs;
				for (int k=0; k<vg.nt; k++) wt[i*vg.nt+k]/=wts;
				for (int k=0; k<vg.nz; k++) wz[i*vg.nz+k]/=wzs;
			}
			
			
			mkdir((string(fnr).substr(0, string(fnr).find_last_of("/"))).c_str(), 0777);
			mkdir((string(fnt).substr(0, string(fnt).find_last_of("/"))).c_str(), 0777);
			mkdir((string(fnz).substr(0, string(fnz).find_last_of("/"))).c_str(), 0777);

			
			// saving
			FILE *fpr = fopen(fnr, "w");
			FILE *fpt = fopen(fnt, "w");
			FILE *fpz = fopen(fnz, "w");
			
			for (int i=0; i<sg.nr; i++) {
				for (int k=0; k<vg.nr; k++) fprintf(fpr, "%u %g\n", i*vg.nr+k, wr[i*vg.nr+k]);
				for (int k=0; k<vg.nt; k++) fprintf(fpt, "%u %g\n", i*vg.nt+k, wt[i*vg.nt+k]);
				for (int k=0; k<vg.nz; k++) fprintf(fpz, "%u %g\n", i*vg.nz+k, wz[i*vg.nz+k]);
			}
			
			
			fclose(fpr);
			fclose(fpt);
			fclose(fpz);
			
			
			delete [] wr;
			delete [] wt;
			delete [] wz;
			
			printf("done\n");
		}
		
		
		void calc_and_save_dm(const char *filename){
			
			printf("Disk::calc_and_save_reldendev('%s')...", filename); fflush(stdout);

			double r1 = 1;
			double r2 = 15;
			uint   nr = 50;
			
			double *m0 = new double[nr];
			double *ms = new double[nr];
			uint *ns = new uint[nr];
			
			memset(ms, 0, nr*sizeof(double));
			memset(m0, 0, nr*sizeof(double));
			memset(ns, 0, nr*sizeof(uint));
		
			
			for (int i=0; i<n; i++) {
				
				double  r = sqrt(SQR(p[i].Pos[0])+SQR(p[i].Pos[1])); 
				double _r = r*(r2-r1)/r2 + r1;
				uint ir = nr*log(_r/r1)/log(r2/r1);
				
				if (ir<nr) ms[ir]+=m/n;
				if (ir<nr) ns[ir]++;
				
			}
			double mp(0);
			double ma(0);
			FILE *fp = fopen(filename, "w");
			for (int i=0; i<nr; i++) {
				
				double _r  = (r1*pow(r2/r1,(double)(i+0.0)/nr)-r1)*r2/(r2-r1);
				double _r_ = (r1*pow(r2/r1,(double)(i+0.5)/nr)-r1)*r2/(r2-r1);
				double  r_ = (r1*pow(r2/r1,(double)(i+1.0)/nr)-r1)*r2/(r2-r1);
				
				double mi = M(_r, 0, r_, 20, 100, 1000);
				
				fprintf(fp, "%i %g %g %g %g %i\n",  i, _r_, ms[i], mi, (ms[i]-mi)/ms[i],  ns[i]);
				ma+=mi;
				mp+=ms[i];
			}
			
			fclose(fp);
			
			delete [] m0;
			delete [] ms;
			delete [] ns;
			
			printf("done\n");

			print();
			
			printf("m  = %.15f\n", m);
			printf("ma = %.15f\n", ma);
			printf("ms = %.15f\n", mp);


			
			//exit(0);
			
			
		}
		
		
		void calc_and_save_zmean(const char *filename){
			
			double r1 = 1;
			double r2 = 15;
			uint   nr = 50;
			
			double *zs = new double[nr];
			uint   *ns = new uint[nr];
			
			memset(zs, 0, nr*sizeof(double));
			memset(ns, 0, nr*sizeof(uint));
			
			for (int i=0; i<n; i++) {
				
				double  r = sqrt(SQR(p[i].Pos[0])+SQR(p[i].Pos[1])); 
				double _r = r*(r2-r1)/r2 + r1;
				uint ir = nr*log(_r/r1)/log(r2/r1);
				
				if (ir<nr) zs[ir]+=fabs(p[i].Pos[2]);
				if (ir<nr) ns[ir]++;
				
			}

			FILE *fp = fopen(filename, "w");
			for (int ir=0; ir<nr; ir++) {
				double _r_ = (r1*pow(r2/r1,(double)(ir+0.5)/nr)-r1)*r2/(r2-r1);
				fprintf(fp, "%i\t%g\t%g\t%i\n",  ir, _r_, zs[ir]/ns[ir], ns[ir]);
			}
			fclose(fp);
			
		}
		
		void center_potmin() {
			
			TreeForceFinder tff;
			tff.G = G;
			tff.NumPart = n;
			tff.force_treeallocate(1.5*n, n);
			tff.force_treebuild(p);
			
			int ci; 
			REAL acc[3];
			double minpot(1e99);
			for(int i=0; i<n; i++) {
				tff.force_treeevaluate(p, p[i].Pos, acc, p[i].pot);
				if (p[i].pot<minpot) {minpot = p[i].pot; ci = i; } 
			}
			tff.force_treefree();
			
			for(int i=0; i<n; i++) 
				*(Vector<REAL>*)p[i].Pos-=*(Vector<REAL>*)p[ci].Pos; 
				
		}
		
		void center_cm() {
			
					
			Vector<REAL> CM; double M(0);
			
			for (uint i=0; i<n; i++) {
		
				CM+=p[i].Mass * (*(Vector<REAL>*)p[i].Pos);
				M+=p[i].Mass;

			}
	
			
			CM/=M;
			CM.print();
			
			for(int i=0; i<n; i++) 
				*(Vector<REAL>*)p[i].Pos-=CM; 
				
		}
		
				
		
		void save_vel(const char *filename) {
			printf("Disk::save_vel_float('%s')...", filename); fflush(stdout);
			FILE *fp = fopen(filename, "wb");
			for (int i=0; i<n; i++) {
				float vel[] = {p[i].Vel[0], p[i].Vel[1], p[i].Vel[2], 1 };
				fwrite(vel, sizeof(float), 4, fp);	
			}
			fclose(fp);
			printf("done\n");
		}
		
		void load_vel(const char *filename) {
			printf("Disk::load_vel_float('%s')...", filename); fflush(stdout);
			FILE *fp = fopen(filename, "rb");
			for (int i=0; i<n; i++) {
				
				float vel[4];
				fread(vel, sizeof(float), 4, fp);
				
				p[i].Vel[0] = vel[0];
				p[i].Vel[1] = vel[1];
				p[i].Vel[2] = vel[2];
				
			}
			fclose(fp);
			printf("done\n");
		}


		void save_pos(const char *filename) {
			printf("Disk::save_pos_float('%s')...", filename); fflush(stdout);
			FILE *fp = fopen(filename, "wb");
			for (int i=0; i<n; i++) {
				float pos[] = {p[i].Pos[0], p[i].Pos[1], p[i].Pos[2], 1 };
				fwrite(pos, sizeof(float), 4, fp);	
			}
			fclose(fp);
			printf("done\n");
		}
		
		
		void load_pos(const char *filename) {
			printf("Disk::load_pos_float('%s')...", filename); fflush(stdout);
			FILE *fp = fopen(filename, "rb");
			for (int i=0; i<n; i++) {
				
				float pos[4];
				fread(pos, sizeof(float), 4, fp);
				
				p[i].Pos[0] = pos[0];
				p[i].Pos[1] = pos[1];
				p[i].Pos[2] = pos[2];
				
			}
			fclose(fp);
			printf("done\n");
		}
		
		
		void free() {
			if (p) delete [] p;
			p = NULL;
		}
		
		void print() {
			
			printf("Disk::print()\n");
			printf(" ├ n = %i\n", n);
			printf(" ├ m = %g\n", m);
			printf(" ├ r0 = %g\n", r0);
			printf(" ├ z0 = %g\n", z0);
			printf(" ├ rmax = %g\n", rmax);
			printf(" ├ zmax = %g\n", zmax);
			printf(" └ G = %g\n\n", G);
				
		}
		
};

#endif
